import os
import io
from PyPDF2 import PdfReader  # 用于读取PDF文件
from PyPDF2.generic import IndirectObject  # 用于处理PDF中的间接对象
from docx import Document  # 用于读取Word文档
from openpyxl import load_workbook  # 用于读取Excel文件
from openpyxl.drawing.image import Image as ExcelImage  # 用于处理Excel中的图片
from pptx import Presentation  # 用于读取PPT文件
from PIL import Image  # 用于图像处理
from utils.file_utils import ensure_directory_exists  # 自定义函数，用于确保目录存在


def extract_images_from_pdf(pdf_path, output_dir):
    """
    从PDF文件中提取图片并保存到指定目录。

    :param pdf_path: PDF文件的路径
    :param output_dir: 保存提取图片的目录
    """
    with open(pdf_path, "rb") as file:
        pdf = PdfReader(file)  # 读取PDF文件
        for page_num in range(len(pdf.pages)):  # 遍历每一页
            page = pdf.pages[page_num]
            xObjects = page["/Resources"].get("/XObject", {})  # 获取资源中的XObject
            for obj_ref in xObjects.values():
                # 获取对象的实际引用
                if isinstance(obj_ref, IndirectObject):
                    obj = pdf.get_object(obj_ref)
                else:
                    obj = obj_ref

                # 检查对象是否为图片
                if "/Subtype" in obj and obj["/Subtype"] == "/Image":
                    size = (obj["/Width"], obj["/Height"])  # 获取图片尺寸
                    data = obj.get_data()  # 获取图片数据
                    mode = "RGB"  # 默认颜色模式
                    # 根据颜色空间设置图片模式
                    if obj["/ColorSpace"] == "/DeviceRGB":
                        mode = "RGB"
                    elif obj["/ColorSpace"] == "/DeviceGray":
                        mode = "L"
                    elif obj["/ColorSpace"] == "/DeviceCMYK":
                        mode = "CMYK"
                    img = Image.frombytes(
                        mode, size, data
                    )  # 将数据转换为Pillow图像对象
                    img.save(
                        os.path.join(
                            output_dir, f"page_{page_num}_img_{id(obj)}.png"
                        )  # 保存图片
                    )


def extract_images_from_word(doc_path, output_dir):
    """
    从Word文档中提取图片并保存到指定目录。

    :param doc_path: Word文档的路径
    :param output_dir: 保存提取图片的目录
    """
    from docx.parts.image import ImagePart

    doc = Document(doc_path)  # 读取Word文档
    for rel in doc.part.rels.values():  # 遍历文档中的关系
        if isinstance(rel.target_part, ImagePart):  # 检查是否为图片部分
            img_data = rel.target_part.image.blob  # 获取图片数据
            img = Image.open(io.BytesIO(img_data))  # 读取图片
            img.save(os.path.join(output_dir, f"image_{rel.rId}.png"))  # 保存图片


def extract_images_from_excel(excel_path, output_dir):
    """
    从Excel文件中提取图片并保存到指定目录。

    :param excel_path: Excel文件的路径
    :param output_dir: 保存提取图片的目录
    """
    # 创建输出目录（如果尚不存在）
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # 加载Excel文件
    wb = load_workbook(
        excel_path, data_only=True  # 使用data_only=True确保加载的是计算后的值
    )

    # 遍历每个工作表
    for sheet_name in wb.sheetnames:
        sheet = wb[sheet_name]

        # 遍历所有图片
        for idx, image in enumerate(sheet._images):
            # 获取图片的数据
            img_data = image.ref  # 获取图片的位置信息
            # 如果图片是嵌入式的，我们需要从Excel中读取实际的图片数据
            if isinstance(image, ExcelImage):
                img_stream = image._data()  # 获取图片的二进制数据流
                img = Image.open(io.BytesIO(img_stream))  # 读取图片
                # 构建文件名并保存图片
                filename = f"{sheet_name}_image_{idx}.png"
                img.save(os.path.join(output_dir, filename))  # 保存图片
                print(f"Saved {filename} to {output_dir}")  # 输出保存信息


def extract_images_from_ppt(ppt_path, output_dir):
    """
    从PPT文件中提取图片并保存到指定目录。

    :param ppt_path: PPT文件的路径
    :param output_dir: 保存提取图片的目录
    """
    prs = Presentation(ppt_path)  # 读取PPT文件
    for slide in prs.slides:  # 遍历每一页幻灯片
        for shape in slide.shapes:  # 遍历幻灯片中的每个形状
            if shape.shape_type == 13:  # 13 是 PPT 中图像的类型码
                img = shape.image  # 获取图片
                img_data = img.blob  # 获取图片数据
                img = Image.open(io.BytesIO(img_data))  # 读取图片
                img.save(
                    os.path.join(
                        output_dir,
                        f"slide_{slide.slide_id}_img_{shape.shape_id}.png",  # 保存图片
                    )
                )


def extract_images_from_file(file_path):
    """
    根据文件类型提取图片并保存到指定目录。

    :param file_path: 文件的路径
    """
    output_dir = os.path.splitext(file_path)[0] + "_images"  # 创建输出目录
    ensure_directory_exists(output_dir)  # 确保输出目录存在

    # 根据文件类型调用不同的提取函数
    if file_path.endswith(".pdf"):
        extract_images_from_pdf(file_path, output_dir)
    elif file_path.endswith((".doc", ".docx")):
        extract_images_from_word(file_path, output_dir)
    elif file_path.endswith((".xls", ".xlsx")):
        extract_images_from_excel(file_path, output_dir)
    elif file_path.endswith((".ppt", ".pptx")):
        extract_images_from_ppt(file_path, output_dir)
    else:
        raise ValueError(f"不支持的文件格式: {file_path}")  # 抛出异常
